#include "device/conv_path.h"

ConvPath::ConvPath(std::shared_ptr<MemInterface> In,
                   std::shared_ptr<MemInterface> Out) {
  InMemPtr = std::move(In);
  OutMemPtr = std::move(Out);
  cvnn_start = false;
  irq_status = false;
  irq_mask = false;
  OutBuffer = new s8[InMemPtr->get_ram_size() * 4];
  lut_port = new Lut;
}

ConvPath::~ConvPath() {
  delete[] OutBuffer;
  delete lut_port;
}

inline s64 get_input(s8 *ifm_start, u32 index_w, u32 index_h, u32 ifm_w,
                     u32 ifm_h, u32 pad_left, u32 pad_right, u32 pad_top,
                     u32 pad_bottom, s8 pad_word) {
  CHECK(index_w < ifm_w + pad_left + pad_right);
  CHECK(index_h < ifm_h + pad_top + pad_bottom);
  if ((index_w < pad_left) || (index_w >= ifm_w + pad_left) ||
      (index_h < pad_top) || (index_h >= ifm_h + pad_top))
    return (s64)pad_word;
  else {
    return (s64)ifm_start[ifm_w * (index_h - pad_top) + index_w - pad_left];
  }
}

void ConvPath::DoConv(s8 *ifm_start, s8 *knl_start, s8 *bias_start,
                      s8 *ofm_start, u32 knl_w) {
  u32 stride_w = kernel_stride & 0xf;
  u32 stride_h = (kernel_stride & 0xf0000) >> 16;
  u32 ifm_w = (ifm_width & 0xfff);
  // u32 ifm_w_line_stride = (ifm_width & 0xf0000) >> 16;
  u32 ifm_h = ifm_height;
  u32 ifm_n = ifm_num;
  u32 ofm_w = ofm_width;
  u32 ofm_h = ofm_height;
  u32 ofm_n = ofm_num;
  u32 pad_left = (padding_idc & 1) ? (padding_size & 0xf) : 0;
  u32 pad_right = (padding_idc & 2) ? ((padding_size & 0xf00) >> 8) : 0;
  u32 pad_top = (padding_idc & 4) ? ((padding_size & 0xf0000) >> 16) : 0;
  u32 pad_bottom = (padding_idc & 8) ? ((padding_size & 0xf000000) >> 24) : 0;
  bool is_depthwise = false;
  bool is_activation = false;
  bool is_leakyrelu = false;
  u8 leaky_lambda = 0;
  is_depthwise = conv_detail & 0x2;
  switch (conv_detail & 0x5) {
  case 0x0:
    break;
  case 0x1:
    is_activation = true;
    break;
  case 0x4:
    is_activation = true;
    is_leakyrelu = true;
    break;
  default:
    CHECK_MSG(false, " conv_detail error");
  }
  if (is_leakyrelu)
    leaky_lambda = (u8)(conv_detail >> 8);
  CHECK(knl_w * knl_w == knl_size);
  CHECK((ifm_w + pad_left + pad_right - knl_w) / stride_w == ofm_w - 1);
  CHECK((ifm_h + pad_top + pad_bottom - knl_w) / stride_h == ofm_h - 1);
  s8 *ifm_tmp = ifm_start;
  s8 *knl_tmp = knl_start;
  s8 *bias_tmp = bias_start;
  s8 *ofm_tmp = ofm_start;
  u32 ifm_real_offset = ifm_offset & 0xffffff;
  u32 knl_real_offset = knl_offset & 0x1ff;
  memset(OutBuffer, 0, ofm_w * ofm_h * ofm_n);
  if (is_depthwise) {
    CHECK(ifm_n == ofm_n);
    for (u32 index_out_n = 0; index_out_n < ofm_n; index_out_n++) {
      for (u32 index_out_h = 0; index_out_h < ofm_h; index_out_h++) {
        for (u32 index_out_w = 0; index_out_w < ofm_w; index_out_w++) {
          s8 *knl_p = knl_tmp;
          u32 index_ifm_w_start = index_out_w * stride_w;
          u32 index_ifm_h_start = index_out_h * stride_h;
          s64 result = 0;
          for (u32 index_knl_h = 0; index_knl_h < knl_w; index_knl_h++) {
            for (u32 index_knl_w = 0; index_knl_w < knl_w; index_knl_w++) {
              result += get_input(ifm_tmp, index_ifm_w_start + index_knl_w,
                                  index_ifm_h_start + index_knl_h, ifm_w, ifm_h,
                                  pad_left, pad_right, pad_top, pad_bottom,
                                  padding_word) *
                        (s64)(*knl_p);
              knl_p++;
            }
          }
          s32 bias_value = *((s32 *)bias_tmp);
          u32 out_scale_value = *((u32 *)(bias_tmp + 4));
          result += bias_value;
          result *= out_scale_value;
          if (out_shift < 0)
            result = result << (-out_shift);
          if (out_shift > 0)
            result = result >> out_shift;
          result = result >> 31;
          result = (result + 1) >> 1;
          if (!lut_bypass)
            result = lut_port->query((s32)result);
         if (is_activation) {
            if (result < 0) {
              if (is_leakyrelu)
                result = (result * leaky_lambda) >> 8;
              else
                result = 0;
            }
          }
          if (result < -127)
            result = -127;
          if (result > 127)
            result = 127;
          ofm_tmp[index_out_h * ofm_w + index_out_w] = (s8)result;
        }
      }
      knl_tmp += knl_real_offset;
      ifm_tmp += ifm_real_offset;
      ofm_tmp += ofm_offset;
      bias_tmp += 8;
    }
  } else {
    for (u32 index_out_n = 0; index_out_n < ofm_n; index_out_n++) {
      for (u32 index_out_h = 0; index_out_h < ofm_h; index_out_h++) {
        for (u32 index_out_w = 0; index_out_w < ofm_w; index_out_w++) {
          s8 *knl_p = knl_tmp;
          s8 *ifm_p = ifm_tmp;
          u32 index_ifm_w_start = index_out_w * stride_w;
          u32 index_ifm_h_start = index_out_h * stride_h;
          s64 result = 0;
          for (u32 index_in_n = 0; index_in_n < ifm_n; index_in_n++) {
            for (u32 index_knl_h = 0; index_knl_h < knl_w; index_knl_h++) {
              for (u32 index_knl_w = 0; index_knl_w < knl_w; index_knl_w++) {
                result += get_input(ifm_p, index_ifm_w_start + index_knl_w,
                                    index_ifm_h_start + index_knl_h, ifm_w,
                                    ifm_h, pad_left, pad_right, pad_top,
                                    pad_bottom, padding_word) *
                          (s64)(*knl_p);
                knl_p++;
              }
            }
            knl_p = knl_p - knl_size + knl_real_offset;
            ifm_p += ifm_real_offset;
          }
          s32 bias_value = *((s32 *)bias_tmp);
          u32 out_scale_value = *((u32 *)(bias_tmp + 4));
          result += bias_value;
          result *= out_scale_value;
          if (out_shift < 0)
            result = result << (-out_shift);
          if (out_shift > 0)
            result = result >> out_shift;
          result = result >> 31;
          result = (result + 1)>>1;
          if (!lut_bypass)
            result = lut_port->query((s32)result);
         if (is_activation) {
            if (result < 0) {
              if (is_leakyrelu)
                result = (result * leaky_lambda) >> 8;
              else
                result = 0;
            }
          }
          if (result < -127)
            result = -127;
          if (result > 127)
            result = 127;
          ofm_tmp[index_out_h * ofm_w + index_out_w] = (s8)result;
        }
      }
      knl_tmp += knl_real_offset * ifm_n;
      ofm_tmp += ofm_offset;
      bias_tmp += 8;
    }
  }
  return;
}

void ConvPath::DoFC(s8 *ifm_start, s8 *knl_start, s8 *bias_start,
                    s8 *ofm_start) {
  bool is_activation = false;
  bool is_leakyrelu = false;
  u8 leaky_lambda = 0;
  switch (fc_detail & 0x3) {
  case 0:
    break;
  case 1:
    is_activation = true;
    break;
  case 2:
    is_activation = true;
    is_leakyrelu = true;
    break;
  default:
    CHECK_MSG(false, "fc_detail error");
  }
  if (is_leakyrelu)
    leaky_lambda = (u8)(fc_detail >> 8);
  s8 *knl_tmp = knl_start;
  s8 *bias_tmp = bias_start;
  s8 *ofm_tmp = ofm_start;
  for (u32 index_output = 0; index_output < fc_output_size; index_output++) {
    s8 *ifm_tmp = ifm_start;
    s64 result = 0;

    for (u32 index_input = 0; index_input < fc_input_size; index_input++) {
      result += (s64)(*ifm_tmp) * (s64)(*knl_tmp);
      ifm_tmp++;
      knl_tmp++;
    }
    s32 bias_value = *((s32 *)bias_tmp);
    u32 out_scale_value = *((u32 *)(bias_tmp + 4));
    result += bias_value;
    result *= out_scale_value;
    if (out_shift < 0)
      result = result << (-out_shift);
    if (out_shift > 0)
      result = result >> out_shift;
    result = result >> 31;
    result = (result + 1) >> 1;
    if (!lut_bypass)
      result = lut_port->query((s32)result);
   if (is_activation) {
      if (result < 0) {
        if (is_leakyrelu)
          result = (result * leaky_lambda) >> 8;
        else
          result = 0;
      }
    }
    if (result < -127)
      result = -127;
    if (result > 127)
      result = 127;
    *ofm_tmp = (s8)result;
    ofm_tmp++;
    bias_tmp += 8;
  }
  return;
}

void ConvPath::Run() {
  if (!cvnn_start)
    return;
  if (irq_status)
    CHECK_MSG(false,"irq_status error");
  u32 valid_ifm_offset = (ifm_offset >> 24) & 0x1f;
  u32 ifm_start_offset = (ifm_start_L << 5) + valid_ifm_offset;
  s8 *ifm_start = reinterpret_cast<s8*>(InMemPtr->get_ram_pointer()) + ifm_start_offset;
  u32 valid_bias_offset = (ofm_bias_start_m >> 16) & 0x1f;
  u32 ofm_bias_offset = (ofm_bias_start_L << 5) + valid_bias_offset;
  s8 *bias_start = reinterpret_cast<s8*>(InMemPtr->get_ram_pointer()) + ofm_bias_offset;
  u32 valid_knl_offset = (knl_offset >> 16) & 0x1f;
  u32 knl_start_offset = (knl_start_L << 5) + valid_knl_offset;
  s8 *knl_start = reinterpret_cast<s8*>(InMemPtr->get_ram_pointer()) + knl_start_offset;
  bool is_fc = false;
  u32 knl_w = 0;
  switch (op_type) {
  case 1:
    knl_w = 1;
    break;
  case 1 << 1:
    knl_w = 3;
    break;
  case 1 << 2:
    knl_w = 5;
    break;
  case 1 << 3:
    knl_w = 7;
    break;
  case 1 << 4:
    knl_w = 9;
    break;
  case 1 << 5:
    knl_w = 11;
    break;
  case 1 << 6:
    is_fc = true;
    break;
  case 1 << 7:
    knl_w = 13;
    break;
  default:
    CHECK_MSG(false, "op_type error");
  }
  if (is_fc)
    DoFC(ifm_start, knl_start, bias_start, OutBuffer);
  else
    DoConv(ifm_start, knl_start, bias_start, OutBuffer, knl_w);
  u64 ofm_start = ((u64)ofm_start_m << 32) + ofm_start_L;
  if (is_fc)
    OutMemPtr->write_mem(ofm_start, OutBuffer, fc_output_size);
  else {
    u64 dram_output_tmp = ofm_start;
    s8 *output_buffer_tmp = OutBuffer;
    for (u32 index = 0; index < ofm_num; index++) {
      OutMemPtr->write_mem(dram_output_tmp, output_buffer_tmp,
                         ofm_width * ofm_height);
      dram_output_tmp += ofm_offset;
      output_buffer_tmp += ofm_width * ofm_height;
    }
  }
  cvnn_start = false;
  irq_status = true;
  return;
}

Lut::Lut() {}

Lut::~Lut() {}

s32 Lut::query(s32 input) {
  CHECK_MSG(false, "LUT not ready");
  return 0;
}
